#include "ImageUndistorter.h"
#include <gvars3/instances.h>
#include <stdlib.h>
#include <dirent.h>
#include <string>
#include <algorithm>
#include <boost/program_options.hpp>
#include <opencv/highgui.h>
#include <opencv/cv.h>

using namespace std;
using namespace CVD;
using namespace GVars3;
namespace config = boost::program_options;

ImageUndistorter::ImageUndistorter()
{
  // First, check if the camera is calibrated.
  // If not, we need to run the calibration widget.
  Vector<NUMTRACKERCAMPARAMETERS> vTest;

  vTest = GV3::get<Vector<NUMTRACKERCAMPARAMETERS> >("Camera.Parameters", ATANCamera::mvDefaultParams, HIDDEN);
  mpCamera = new ATANCamera("Camera");
  if(vTest == ATANCamera::mvDefaultParams)
  {
    cout << endl;
    cout << "! Camera.Parameters is not set, need to run the CameraCalibrator tool" << endl;
    cout << "  and/or put the Camera.Parameters= line into the appropriate .cfg file." << endl;
    exit(1);
  }

  gvar3<int> gvnCameraWidth("Camera.Width", 640, GVars3::HIDDEN | GVars3::FATAL_IF_NOT_DEFINED); 
  gvar3<int> gvnCameraHeight("Camera.Height", 480, GVars3::HIDDEN | GVars3::FATAL_IF_NOT_DEFINED); 

  mirUndistortSize = ImageRef(*gvnCameraWidth, *gvnCameraHeight);
  mimUndistortMap.resize(mirUndistortSize);
  mimUndistortedRGB.resize(mirUndistortSize);

   for(int i = 0; i < mirUndistortSize.x; ++i)  
    for(int j = 0; j < mirUndistortSize.y; ++j)
      mimUndistortMap[j][i] = mpCamera->UndistortMap(makeVector(i,j));
}

ImageUndistorter::~ImageUndistorter()
{
  delete mpCamera;
}


const CVD::BasicImage<CVD::Rgb<CVD::byte> >& ImageUndistorter::Undistort(const CVD::BasicImage<CVD::Rgb<CVD::byte> > &imDistortedRGB)
{
  ImageRef irDistortedSize = imDistortedRGB.size();
  if(mirUndistortSize == irDistortedSize)
  {
    for(int i = 0; i < mirUndistortSize.x; ++i)  
    {
      for(int j = 0; j < mirUndistortSize.y; ++j)
      {
        mimUndistortedRGB[j][i] = BilinearPixel(imDistortedRGB, mimUndistortMap[j][i]);
      }
    }
    return mimUndistortedRGB;
  }
  else
  {
    std::cout << "Warning: Size mismatch between image and camera parameters, skipping undistortion" << std::endl;
    return imDistortedRGB;
  }
}

void StringToDouble(const string str, double &number)
{
  stringstream buff;

  buff << str;
  buff >> number;
}


void DoubleToString(double number, string &str)
{
  stringstream buff;

  buff << number;
  buff >> str;
}


void GetImageList(string dir_name, string file_prefix, string file_type, vector<string>& file_list)
{
  DIR *dp;
  struct dirent *dirp;
  string filename;
  string extension;
  vector<double> frame_number;
  double number;
  unsigned int extension_pos, prefix_pos;
  int prefix_length = file_prefix.length();
	bool leading_zeros = false;
	int number_digits = 0;
  stringstream buff;

  if((dp  = opendir(dir_name.c_str())) == NULL) {
    cout << "Error: Unable to open directory " << dir_name << endl;
    exit(1);
  }

  while ((dirp = readdir(dp)) != NULL) {
    filename = string(dirp->d_name);
    extension_pos = filename.find(".");
    prefix_pos = filename.find(file_prefix);
    if(extension_pos != string::npos && prefix_pos != string::npos)
    {
      extension = filename.substr(extension_pos+1, string::npos);

			if(extension == file_type)
			{
		    filename.erase(extension_pos, string::npos); //erase extension
		    filename.erase(0, prefix_length); //erase 'Frame'
				if(filename.compare(0,1,"0") == 0 && filename.length() > 1)
				{
					leading_zeros = true;
					number_digits = filename.length();
				}
		    StringToDouble(filename, number);
		    frame_number.push_back(number);
			}
      else
        std::cout << "Warning: Ignored file with invalid extension (" << filename  << " " << extension << ")" << std::endl;
    }
  }
  closedir(dp);

  sort(frame_number.begin(), frame_number.end());

  for(vector<double>::iterator iter = frame_number.begin(); iter < frame_number.end(); iter++)
  {
    DoubleToString(*iter, filename);
    //if(leading_zeros)
      //std::cout << number_digits-filename.length() << std::endl;
		if(leading_zeros)
			filename.insert(0, number_digits-filename.length(), '0');
    filename.insert(0, file_prefix);
    filename.insert(0, dir_name);
    filename.append(".").append(extension);
    file_list.push_back(filename);
  }
}


int main(int argc, char* argv[])
{
  cout << "  Welcome to CameraCalibrator " << endl;
  cout << "  -------------------------------------- " << endl;
  cout << "  Parallel tracking and mapping for Small AR workspaces" << endl;
  cout << "  Copyright (C) Isis Innovation Limited 2008 " << endl;
  cout << endl;  

  GUI.LoadFile("./config/settings.cfg");
  GUI.StartParserThread();
  atexit(GUI.StopParserThread); // Clean up readline when program quits

	// Setup program options
  string directory, prefix, extension;
  bool verbose;
	config::options_description desc("Allowed options");
	desc.add_options()
									("help", "produce help message")
									("directory,d", config::value<string>(&directory)->default_value(""),	"directory containing images to undistort")
									("prefix,p", config::value<string>(&prefix)->default_value(""),	"image prefix, ie. image_ for image_000, image_001 etc.")
									("extension,e", config::value<string>(&extension)->default_value(""),	"image extension, ie. jpg/jpeg")
									("verbose,v", config::bool_switch(&verbose)->default_value(false),"verbose output");
	config::variables_map vm;
	config::store(config::parse_command_line(argc, argv, desc), vm);
	config::notify(vm);

	if (vm.count("help"))
  {
		std::cout << desc << std::endl;
		return 1;
	}
  
  ImageUndistorter undist;

  vector<string> images;
  GetImageList(directory, prefix, extension, images);

  for(int i = 0; i < images.size(); ++i)
  {
    std::cout << images[i] << std::endl;
    IplImage *img = cvLoadImage(images[i].c_str());
    BasicImage<CVD::Rgb<CVD::byte> > bimg((CVD::Rgb<CVD::byte>*)(img->imageData), ImageRef(img->width, img->height));

    CVD::Image<CVD::Rgb<CVD::byte> > imDistortedRGB;
    imDistortedRGB.resize(ImageRef(img->width, img->height));
    imDistortedRGB.copy_from(bimg);

    const CVD::BasicImage<CVD::Rgb<CVD::byte> > &imUndistorted = undist.Undistort(imDistortedRGB);
    IplImage* img_out = cvCreateImageHeader(cvGetSize(img), img->depth, img->nChannels);    
    cvSetData(img_out, (void*)imUndistorted.data(), img->widthStep);

    int prefix_pos = images[i].find(prefix);
    string filename = images[i].insert(prefix_pos, "undistorted_");
    int extension_pos = filename.find_last_of(".");
    filename.replace(extension_pos, string::npos, ".jpg");
std::cout << filename << std::endl;

    cvSaveImage(filename.c_str(), img_out);
    cvReleaseImage(&img);
    cvReleaseImageHeader(&img_out);
  }
}

